package com.chs.springai_alibaba.service;

import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingModel;
import com.alibaba.cloud.ai.dashscope.rag.*;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentRetriever;
import org.springframework.stereotype.Service;

import java.util.Arrays;
import java.util.List;

@Service
public class RAGService {
    private final DashScopeApi dashscopeApi;
    public RAGService(DashScopeApi dashscopeApi) {
        this.dashscopeApi = dashscopeApi;
    }

    public DocumentRetriever importDocuments(){
        // 1. 解析文档和chunk切分
        String filePath = "C:\\Users\\kubayaxi\\Desktop\\haha.txt";
        DashScopeDocumentCloudReader reader = new DashScopeDocumentCloudReader(filePath, dashscopeApi, null);
        List<Document> documentList = reader.get();
        DashScopeDocumentTransformer transformer = new DashScopeDocumentTransformer(dashscopeApi);
        List<Document> transformerList = transformer.apply(documentList);
        System.out.println(transformerList.size());

        // 2. 文档向量化
        DashScopeEmbeddingModel embeddingModel = new DashScopeEmbeddingModel(dashscopeApi);
        Document document = new Document("你好阿里云");
        float[] vectorList = embeddingModel.embed(document);

        // 3. 导入文档内容到向量存储
        DashScopeCloudStore cloudStore = new DashScopeCloudStore(dashscopeApi, new DashScopeStoreOptions("subway"));

        cloudStore.add(Arrays.asList(document));

        // 4. 删除文档
        cloudStore.delete(Arrays.asList(document.getId()));

        DocumentRetriever retriever = new DashScopeDocumentRetriever(dashscopeApi, DashScopeDocumentRetrieverOptions.builder()
                .withIndexName("subway")
                .build());
        return retriever;
    }

}
